package com.leozhang.utils;

import java.io.ByteArrayOutputStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import javax.crypto.Cipher;
import org.apache.commons.codec.binary.Base64;

public class RsaUtil {
    private static final String ALGORITHM_RSA = "RSA";

    private static final String ALGORITHM_SIGNATURE = "SHA1WithRSA";

    public static final int KEY_SIZE = 2048;

    private static final String PADDING = "RSA/ECB/PKCS1PADDING";

    public static String[] generateRSAKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        PublicKey publicKey = keyPair.getPublic();
        PrivateKey privateKey = keyPair.getPrivate();
        String publicKeyStr = Base64.encodeBase64String(publicKey.getEncoded());
        String privateKeyStr = Base64.encodeBase64String(privateKey.getEncoded());
        return new String[] { publicKeyStr, privateKeyStr };
    }

    public static String encryptByPublicKeyStr(String publicKeyStr, String data) throws Exception {
        PublicKey publicKey = getPublicKeyFromString(publicKeyStr);
        byte[] bytes = doCipher(1, publicKey, data.getBytes("utf-8"));
        return Base64.encodeBase64String(bytes);
    }

    public static String decryptByPrivateKeyStr(String privateKeyStr, String encryptedData) throws Exception {
        PrivateKey privateKey = getPrivateKeyFromString(privateKeyStr);
        byte[] bytes = doCipher(2, privateKey, Base64.decodeBase64(encryptedData));
        return new String(bytes, "utf-8");
    }

    public static String encryptByPublicKeyFile(String publicKeyPath, String data) throws Exception {
        PublicKey publicKey = getPublicKeyFromFile(publicKeyPath);
        byte[] bytes = doCipher(1, publicKey, data.getBytes("utf-8"));
        return Base64.encodeBase64String(bytes);
    }

    public static String decryptByPrivateKeyFile(String privateKeyPath, String encryptedData) throws Exception {
        PrivateKey privateKey = getPrivateKeyFromFile(privateKeyPath);
        byte[] bytes = doCipher(2, privateKey, Base64.decodeBase64(encryptedData));
        return new String(bytes, "utf-8");
    }

    public static String encryptByPriKeyFile(String privateKeyPath, String data) throws Exception {
        PrivateKey privateKey = getPrivateKeyFromFile(privateKeyPath);
        byte[] bytes = doCipher(1, privateKey, data.getBytes("utf-8"));
        return Base64.encodeBase64String(bytes);
    }

    public static String decryptByPublicKeyFile(String publicKeyPath, String encryptedData) throws Exception {
        PublicKey publicKey = getPublicKeyFromFile(publicKeyPath);
        byte[] bytes = doCipher(2, publicKey, Base64.decodeBase64(encryptedData));
        return new String(bytes, "utf-8");
    }

    private static PublicKey getPublicKeyFromString(String publicKeyStr) throws Exception {
        byte[] encodedKey = Base64.decodeBase64(publicKeyStr);
        return generatePublic(encodedKey);
    }

    private static PublicKey getPublicKeyFromFile(String publicKeyPath) throws Exception {
        byte[] encodedKey = Files.readAllBytes(Paths.get(publicKeyPath, new String[0]));
        return generatePublic(encodedKey);
    }

    private static PublicKey generatePublic(byte[] encodedKey) throws Exception {
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(encodedKey);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        PublicKey publicKey = keyFactory.generatePublic(keySpec);
        return publicKey;
    }

    private static PrivateKey getPrivateKeyFromString(String privateKeyStr) throws Exception {
        byte[] encodedKey = Base64.decodeBase64(privateKeyStr);
        return generatePrivate(encodedKey);
    }

    private static PrivateKey getPrivateKeyFromFile(String privateKeyPath) throws Exception {
        byte[] encodedKey = Files.readAllBytes(Paths.get(privateKeyPath, new String[0]));
        return generatePrivate(encodedKey);
    }

    private static PrivateKey generatePrivate(byte[] encodedKey) throws Exception {
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(encodedKey);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        PrivateKey privateKey = keyFactory.generatePrivate(keySpec);
        return privateKey;
    }

    private static byte[] doCipher(int mode, Key key, byte[] bytes) throws Exception {
        int segmentSize;
        if (mode == 2) {
            segmentSize = 256;
        } else {
            segmentSize = 245;
        }
        Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1PADDING");
        cipher.init(mode, key);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int inputLen = bytes.length;
        int offSet = 0;
        int i = 0;
        while (inputLen - offSet > 0) {
            byte[] cache;
            if (inputLen - offSet > segmentSize) {
                cache = cipher.doFinal(bytes, offSet, segmentSize);
            } else {
                cache = cipher.doFinal(bytes, offSet, inputLen - offSet);
            }
            out.write(cache, 0, cache.length);
            i++;
            offSet = i * segmentSize;
        }
        byte[] data = out.toByteArray();
        out.close();
        return data;
    }

    private static byte[] sign(PrivateKey privateKey, byte[] data) throws Exception {
        Signature signature = Signature.getInstance("SHA1WithRSA");
        signature.initSign(privateKey);
        signature.update(data);
        return signature.sign();
    }

    private static boolean verify(PublicKey publicKey, byte[] data, byte[] sign) throws Exception {
        Signature signature = Signature.getInstance("SHA1WithRSA");
        signature.initVerify(publicKey);
        signature.update(data);
        return signature.verify(sign);
    }
}
